import torch
from torch.nn.modules import loss
from dataset import AttentionDataset
from torch.utils.data import DataLoader
import os
from utils import ANNAVQA


os.environ['CUDA_VISIBLE_DEVICES'] = '8'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
# raise os.error

train_dataset = AttentionDataset('test.pkl')
# train_dataset = AttentionDataset('train.pkl')
test_dataset = AttentionDataset('test.pkl')
train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=8, shuffle=True)

# total samples of data and number of iterations performed
total_samples = len(train_dataset)
n_iterations = total_samples//4
print(total_samples, n_iterations)
# for i, (targets, labels) in enumerate(train_dataloader):
#     print(targets, labels)


def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y, length, seg_video_len) in enumerate(dataloader):
        print(f"batch: {batch}")
        X, y = X.to(device), y.to(device)
        # Compute prediction and loss
        # print('X', X.shape)  # [2, 48, 4096]
        # print('y', y.shape)  # [2]
        # print('length', length)  # [2, 1, 1] -> 48
        print('video_length', seg_video_len)  # [2,1,1] -> 24
        # raise os.error
        y = y.to(torch.float32)
        print(f"batch: {batch}.p1")
        pred = model(X, length, 24)
        print(f"batch: {batch}.p2")
        loss = loss_fn(pred, y)
        print(f"batch: {batch}.p3")

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for batch, (X, y, length, seg_video_len) in enumerate(dataloader):
            pred = model(X, length, 24)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


if __name__ == '__main__':
    print('load and train')
    model = ANNAVQA()
    # model = torch.nn.DataParallel(model) # multi-GPU
    # model = torch.nn.DataParallel(model,device_ids=[0,1]) # multi-GPU
    loss_fn = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    # 定义模型
    model.to(device) # 移动模型到cuda

    epochs = 1
    for t in range(epochs):
        print(f"Epoch {t+1}\n-------------------------------")
        train_loop(train_dataloader, model, loss_fn, optimizer)
        # test_loop(test_dataloader, model, loss_fn)
    torch.save(model, 'model.pth')
    print("Done!")
